// Copyright (c) Microsoft Corporation. All rights reserved.

#include <windows.h>
#include <strsafe.h>

#ifdef _DEBUG
#define ODS(m) OutputDebugStringW(L"CodeRunnerHelper.dll: " m)
#else
#define ODS(m) __noop(m)
#endif

#define MAX_MESSAGE_LENGTH 32767
#define COUNT_OF(x) (sizeof(x)/sizeof(x[0]))
#define LENGTH_OF(x) (COUNT_OF(x) - 1)

void ProcessRequest(wchar_t *pRequest, int requestLength, wchar_t *pReply, int *pReplyLength, BOOL *pDisconnect)
{
	LCID lcid = MAKELCID(MAKELANGID(LANG_ENGLISH, SUBLANG_ENGLISH_US), SORT_DEFAULT);
	wchar_t const setCommand[] = L"set\r\n";
	wchar_t const terminateCommand[] = L"terminate";

	*pReplyLength = 0;

	if (requestLength >= LENGTH_OF(setCommand) && 
		CompareString(lcid, 0, pRequest, LENGTH_OF(setCommand), setCommand, LENGTH_OF(setCommand)) == CSTR_EQUAL)
	{
		// Message format:
		// 
		//   set
		//   name=value
		//

		// Is the request too short to proceed further?
		if (requestLength <= 6)  // set\r\n\0
			return;

		wchar_t *pName = pRequest + LENGTH_OF(setCommand);
		wchar_t *pValue = pName + 1;
	
		while (*pValue && (pValue - pRequest) < requestLength - 2 && *pValue != L'=')
			pValue++;

		// TODO: Allow setting multiple environment variables at once?
			
		if (*pValue == L'=')
		{
			*pValue++ = 0;
		
			SetEnvironmentVariableW(pName, pValue);	
		}
	}
	else if (requestLength >= LENGTH_OF(terminateCommand) && 
		CompareString(lcid, 0, pRequest, LENGTH_OF(terminateCommand), terminateCommand, LENGTH_OF(terminateCommand)) == CSTR_EQUAL)
	{
		// Message format:
		//
		//   terminate
		//
		
		*pDisconnect = TRUE;
	}
}

// This method is called from the remotely injected thread
DWORD RealControlRemoteEnvironment(void)
{
	LPTSTR pipeName = L"\\\\.\\pipe\\coderunner$"; 
	HANDLE hPipe = CreateFile(pipeName, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);

	if (hPipe == INVALID_HANDLE_VALUE) 
	{
		// If we don't get connected then we just bail, because obviously things aren't well in the outside world
		return 0;
	}

	wchar_t pRequest[MAX_MESSAGE_LENGTH]; 
	wchar_t pReply[MAX_MESSAGE_LENGTH]; 
	DWORD actualBytes; 
	BOOL success; 
	BOOL disconnect = FALSE;
 
	while (!disconnect) 
	{ 
		int attemptedBytes = MAX_MESSAGE_LENGTH * sizeof(wchar_t);

		success = ReadFile(hPipe, pRequest, attemptedBytes, &actualBytes, NULL);

		if (!success || actualBytes == 0 || actualBytes % sizeof(wchar_t) != 0) 
			break; 
		
		int requestLength = (actualBytes - 1) / sizeof(wchar_t);
		
		// Check for null terminated string
		if (pRequest[requestLength] != 0)
			break;
		
		int replyLength;

		ProcessRequest(pRequest, requestLength, pReply, &replyLength, &disconnect);
		
		if (replyLength > 0)
		{
			attemptedBytes = (replyLength + 1) * sizeof(wchar_t);
			
			success = WriteFile(hPipe, pReply, attemptedBytes, &actualBytes, NULL);
			
			if (!success || actualBytes != attemptedBytes) 
				break;
		}
	}

	// Flush the pipe to allow the client to read the pipe's contents before disconnecting. Then disconnect 
	// the pipe, and close the handle to this pipe instance. 
	FlushFileBuffers(hPipe); 
	DisconnectNamedPipe(hPipe); 
	CloseHandle(hPipe); 

	return 0;
}

#ifndef PIPESERVER

#pragma comment(linker, "/EXPORT:ControlRemoteEnvironment=_ControlRemoteEnvironment@0")
extern "C" DWORD WINAPI ControlRemoteEnvironment(void)
{
	DWORD ret = 0;

	__try
	{
		ret = RealControlRemoteEnvironment();
	}
	__except(EXCEPTION_CONTINUE_EXECUTION)
	{
		// We MUST NOT bring down the process into which we were injected!
	}

	return ret;	
}

#endif